from rest_framework.permissions import IsAuthenticated
from rest_framework.response import Response
from rest_framework.views import APIView

from six_it.models import Enterprise


class CollectionEnterpriseView(APIView):
    '''收藏公司'''
    permission_classes = [IsAuthenticated]

    def post(self, request, id):
        enterprise = Enterprise.objects.get(id=id)
        user = request.user
        enterprise.users.add(user)
        enterprise.save()
        return Response({
            "message": "更新成功",
            "success": True
        })


class CancelCollectEnterpriseView(APIView):
    '''取消收藏公司'''
    permission_classes = [IsAuthenticated]

    def post(self, request, id):
        enterprise = Enterprise.objects.get(id=id)
        user = request.user
        enterprise.users.remove(user)
        enterprise.save()
        return Response({
            "message": "取消收藏成功",
            "success": True
        })
